#include <stdio.h>
#include <stdlib.h> 
#include <time.h>
#include "rsa.h"

int main()
{
    uint64_t n;
    uint64_t e;
    uint64_t d;
    keygen(&n, &e, &d);
    uint32_t m = 20215521;
    printf("原文:%lu\n", m);
    uint64_t c = encrypt(m, e, n);
    printf("密文:%llu\n", c);
    uint64_t m_ = decrypt(c, d, n);
    printf("译文:%llu\n", m_);
    /*
    uint32_t m = 19;
    printf("原文:%lu\n", m);
    uint64_t c = encrypt(m, 5, 119);
    printf("密文:%llu\n", c);
    uint32_t m_ = decrypt(c, 77, 119);
    printf("译文:%lu\n", m_);
    */
}

void keygen(uint64_t *n, uint64_t *e, uint64_t *d)
{
    // 选取两个保密的大素数p和q（二者不可相差悬殊）
    uint32_t p = big_prime();
    uint32_t q = big_prime();
    printf("p:%lu\n", p);
    printf("q:%lu\n", q);
    // 计算 n = pq，φ(n) = (p-1)(q-1)
    *n = (uint64_t) p * q;
    uint64_t phi = (uint64_t) (p - 1) * (q - 1);
    printf("n:%llu\n", *n);
    printf("φ:%llu\n", phi);
    // 选取e，使得 1 < e < φ(n) 而且 gcd(e, φ(n)) = 1 
    while(true)
    {
        random(e, sizeof(uint64_t));
        if (1 < *e && *e < phi && gcd(*e, phi) == 1)
            break;
    }
    printf("e:%llu\n", *e);
    // 计算d，使得 de ≡ 1 mod φ(n)
    *d = mod_mul_inv(*e, phi);
    printf("d:%llu\n", *d);
}

uint32_t big_prime() {
    uint32_t p;
    while(true)
    {
        random(&p, sizeof(uint32_t));
        if (UINT16_MAX < p && miller_rabin_is_prime(p))
            break;
    }
    return p;
}

void random(void *bytes, int num)
{
    static bool initialized = false;  
    if (!initialized) {
        srand((unsigned int)time(NULL));
        initialized = true;
    }
    char *arr = (char *) bytes;
    for (size_t i = 0; i < num; i++)
    {
        arr[i] = rand();
    }
}


bool miller_rabin_is_prime(uint64_t n)
{
    // 参考https://zhuanlan.zhihu.com/p/349360074 
    // 如果n是偶数，则可以直接特判
    if (n < 2)
        return false;
    if (n == 2)
        return true;
    if (n % 2 == 0)
        return false;
    // 7轮检验
    static const uint64_t A[] = { 2, 325, 9375, 28178, 450775, 9780504, 1795265022 };
    for (size_t i = 0; i < sizeof(A); i++)
    {
        // 如果n与a不互素，则不能用这个a
        uint64_t gcd_n_a = gcd(n, A[i]);
        if (1 < gcd_n_a)
        {
            if (gcd_n_a < n)
                return false;
            continue;
        }
        
        if (!miller_rabin_test(n, A[i]))
            return false;
    }
    return true;
}

bool miller_rabin_test(uint64_t n, uint64_t a)
{
    // 费马小定理：对于素数p，若a与p互素，则 a^(p-1) ≡ 1 (mod p)
    if (mod_exp(a, n - 1, n) > 1)
        return false;
    // 二次探测定理：对于素数p，若 x^2 ≡ 1 (mod p)，则小于p的解只有两个，1 或 p-1
    for (uint64_t e = n - 1; e % 2 == 0 ; e /= 2)
    {
        uint64_t r = mod_exp(a, e, n);
        if (r == 1)
            continue;
        else if (r == n - 1)
            return true;
        else
            return false;
    }
    return true;
    /* 上述代码也可以写成下面这样：
    // 设 n-1 = u * 2^t （u是奇数）
    uint64_t u = n - 1;
    int t = 0;
    while (u % 2 == 0)
    {
        u /= 2;
        t++;
    }
    // 检验 a^u，a^2u，a^4u，a^8u ... a^(u * 2^(t-1))
    uint64_t v = mod_exp(a, u, n);
    for (size_t i = 0; i < t; i++)
    {
        if (v == 1)
        {
            return true; // 1是不动点，所以可以直接return true
        }
        else if (v == n - 1)
        {
            v = mod_mul(v, v, n);
            if (v == 1)
                return true;
            else
                return false; // 如果不会违反费马小定理，那么将会违反二次探测定理
        }
        else
        {
            v = mod_mul(v, v, n);
            if (v == 1)
                return false; // 违反二次探测定理
        }
    }
    return false; // 违反费马小定理
    */
}

uint64_t mod_exp(uint64_t b, uint64_t e, uint64_t m)
{
    // 快速幂算法
    uint64_t result = 1;
    uint64_t v = b;
    for (uint64_t bit = 1; bit != 0; bit <<= 1)
    {
        if (e & bit)
        {
            result = mod_mul(result, v, m);
        }
        v = mod_mul(v, v, m);
    }
    return result;
}

uint64_t mod_mul(uint64_t a, uint64_t b, uint64_t m)
{
    uint64_t result = 0;
    uint64_t v = b;
    for (uint64_t bit = 1; bit != 0; bit <<= 1)
    {
        if (a & bit)
        {
            result = mod_add(result, v, m);
        }
        v = mod_add(v, v, m);
    }
    return result;
}

uint64_t mod_add(uint64_t a, uint64_t b, uint64_t m)
{
    // 为了防止溢出，不能简单直接return (a + b) % m;
    a %= m;
    b %= m;
    if (a >= m - b)
        return a - (m - b);
    else
        return a + b;
}

uint64_t gcd(uint64_t a, uint64_t b)
{
    // 参考https://zhuanlan.zhihu.com/p/100567253
    if (b == 0)
        return a;
    return gcd(b, a % b);
}

uint64_t extended_gcd(uint64_t a, uint64_t b, int64_t *x, int64_t *y)
{
    if (b == 0)
    {
        *x = 1;
        *y = 0;
        return a;
    }
    // a = k*b + c
    // c = a - k*b
    uint64_t k = a / b;
    uint64_t c = a % b;
    int64_t u;
    int64_t v;
    uint64_t g = extended_gcd(b, c, &u, &v);
    // u*b + v*c = g
    // u*b + v*(a - k*b) = g
    // u*b + v*a - v*k*b = g
    // v*a + (u - v*k)*b = g
    *x = v;
    *y = u - v*k;
    return g;
}

// 上述代码中的x和y有可能溢出，因此下面代码使用符号位，避免溢出
// 1代表正数，0代表为零，-1代表负数
uint64_t __extended_gcd(uint64_t a, uint64_t b, uint64_t *x, int8_t *sign_x, uint64_t *y, int8_t *sign_y)
{
    if (b == 0)
    {
        *x = 1;
        *sign_x = +1;
        *y = 0;
        *sign_y = 0;
        return a;
    }
    // a = k*b + c
    // c = a - k*b
    uint64_t k = a / b;
    uint64_t c = a % b;
    uint64_t u;
    uint64_t v;
    int8_t sign_u;
    int8_t sign_v;
    uint64_t g = __extended_gcd(b, c, &u, &sign_u, &v, &sign_v);
    // u*b + v*c = g
    // u*b + v*(a - k*b) = g
    // u*b + v*a - v*k*b = g
    // v*a + (u - v*k)*b = g
    *x = v;
    *sign_x = sign_v;
    if (sign_u > 0) // 并且 sign_v <= 0
    {
        *y = u + v*k;
        *sign_y = +1;
    }
    else if (sign_u == 0)
    {
        *y = v*k;
        *sign_y = -sign_v;
    }
    else // sign_u < 0 并且 sgin_v >= 0
    {
        *y = u + v*k;
        *sign_y = -1;
    }
    return g;
    // 拓展欧几里得算法有个特点：|x| <= b 并且 |y| <= a 并且x和y的符号相反（或者其一为0）
    // 这个特点保证了x和y不会溢出
    // 可以使用数学归纳法证明上述特点：
    // (I)当 0 <= u <= a%b 且 -b <= v <= 0时，
    // 因为x = v，所以 -b <= x <= 0
    // 因为y = u - v*k，所以 0 -0*k <= y <= a%b -(-b)*k，所以 0 <= y <= a
    // (II)当 0 <= v <= b 且 -a%b <= u <= 0时，
    // 因为x = v，所以 0 <= x <= b
    // 因为y = u - v*k，所以 -a%b - b*k <= y <= 0 - 0*k，所以 -a <= y <= 0
    // 综上所述
    // 因为条件(I)导致条件(II)，条件(II)导致条件(I)，所以数学归纳法的递推得到满足
    // 因为x = 1，y = 0满足条件(I)，所以数学归纳法的基础得到满足
}

uint64_t mod_mul_inv(uint64_t a, uint64_t b)
{
    // 计算x，使得 ax ≡ 1 mod b
    // 即 ax = 1 + kb (k ∈ Z)
    // 即 ax - bk = 1
    // 即 ax + by = 1，可用拓展欧几里得算法求解
    uint64_t x;
    int8_t sign_x;
    uint64_t y;
    int8_t sign_y;
    __extended_gcd(a, b, &x, &sign_x, &y, &sign_y);
    if (sign_x >= 0)
    {
        return x;
    }
    else
    {
        // ax + by = 1
        // a(x + b) + b(y - a) = 1
        return b - x;
    }
}

uint64_t encrypt(uint64_t m, uint64_t e, uint64_t n)
{
    return mod_exp(m, e, n);
}

uint64_t decrypt(uint64_t c, uint64_t d, uint64_t n)
{
    return mod_exp(c, d, n);
}
